'''
Author: Zhuangyu
Date: 2022-06-02 14:31:11
LastEditTime: 2022-06-13 23:54:57
'''
from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from Parameters import *
from scipy import integrate, linalg, interpolate
from tqdm import tqdm
import mpctools as mpc
import casadi
import scipy.io as sio  
import van_Genuchten as vg

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
DeltaT,Nsim=time_parameters()
p=Loam() #Soil parameters

Na=Nz
Nu=1

def RichardsModel_1D(x,u,ui):
    
    #x is the state vector
    #u is the irrigation amount (the input)
    #w is the process noise
    #Kc is the crop coefficient
    #ET0 is the reference evapotranspiration
    #p represents the soil parameters
    
    Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
    # dhdt=SX.zeros(Nz)
    dhdt=np.zeros(Nz)
    
    Kc=0.88      #Crop coefficient

    ET0=1.389e-08 #Reference evapotranspiration
    
    for i in range(0,Nz):

        CurrentState=x[i]         
                
        if i == 0: #Bottom boundary, Free drainage boundary condition(BC)
            BCz_L=CurrentState
            BCz_U=x[i+1]
         
        elif i == Nz-1: #Top boundary
            BCz_L = x[i - 1]    
         
        else: #Internal nodes
            BCz_L = x[i - 1]
            BCz_U = x[i + 1]

        #Computing the unsaturated hydraulic conductivity at the centre of the compartments
        KzL=0.5*(vg.KFun(CurrentState,p)+vg.KFun(BCz_L,p))
        KzU=0.5*(vg.KFun(CurrentState,p)+vg.KFun(BCz_U,p))
        
        #Computing the pressure head gradient
        DHzL=(CurrentState-BCz_L)/dz 
        DHzU=(BCz_U-CurrentState)/dz

        if i == Nz-1: #Surface Node [Nuemann BC which incorporates the irrigation amount, u]
            Term1 = (1.0/dz)*( -u - KzL*DHzL)
            Term2 =(1.0/dz)*( -1*KzL)
        
        else: #Nodes below the surface
            Term1 = (1.0/dz) * (KzU*DHzU - KzL*DHzL)
            Term2 = (1.0/dz) * (KzU - KzL)

        Term3 = (Kc*ET0)/Hz #The Sink term [Root water extraction rate]
        Term4 = Term1 + Term2 - Term3
        Term5 = Term4/vg.CFun(CurrentState,p)

        dhdt[i]=Term5
    DHDT=dhdt + ui
    return DHDT

# F = mpc.getCasadiFunc(RichardsModel_1D,[Nz,Nu,Na],["x","u","a"],"RichardsModel_1D")
# F_rk4 = mpc.getCasadiFunc(RichardsModel_1D,[Nz,Nu,Na],["x","u","a"],
#                                    "ode_rk4",rk4=True,Delta=DeltaT,M=1)
# Richards = mpc.DiscreteSimulator(RichardsModel_1D, DeltaT, [Nz,Nu,Na], ["x","u","a"])

def getOutputs1D_allnodes(x,p):
    #Returns the measurements for all the states
    Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
    y = vg.thetaFun(x,p)
    CMatrix = cmatrix_single(Nz)
    # y_pred = mtimes(CMatrix,y)
    y_pred = np.dot(CMatrix,y)
    return y_pred

def getOutputs1D(C,x,p):
    #Returns measurements for some selected states based on the C-matrix
    y = vg.thetaFun(x,p)
    # y_pred = mtimes(C,y)
    y_pred = np.dot(C,y)    
    return y_pred

u=np.zeros(6*120)
uu=np.zeros(Nsim)
for i in range(6*120):
    if i < 8:
        u[i]=-5e-06 #Moisture is supplied only at the first time instant
    else:
        u[i]=0
uu=np.tile(u,(1,60))
uu=np.transpose(uu)

#Initial condition [pressure head]
x0=-0.8*np.ones(Nz)

#Numpy array to store the simulation results [pressure head]
xx=np.zeros((Nsim+1,Nz))
xx[0,:]=x0

ode1=np.zeros((Nsim+1,Nz))

#Model uncertainty
np.random.seed(2)
w=0.00001*np.random.randn(Nsim,Nz)
# w=0.01*np.random.randn(Nsim,Nz)
np.random.seed(3)
vv=0.01*np.random.randn(Nsim+1,Nz) #measurement noise

#Numpy array to store the simulation results [volumetric moisture content]
yy=np.zeros((Nsim+1,Nz))
yy[0,:]=getOutputs1D_allnodes(x0,p).ravel()+vv[0,:]

#unknown inputs
a1=0*np.ones((Nsim+1,Nz))
a=0.0000005*np.ones((Nsim+1,Nz))

e0 = time.time()

for i in tqdm(range(1, Nsim+1)):
    # Generate observed data (x bias)
    ode1[i-1,:]= RichardsModel_1D(tuple(xx[i-1]), uu[i-1],a[i-1])
    xx[i,:]=xx[i-1,:]+DeltaT*ode1[i-1,:]+w[i-1,:]
    yy[i,:]=getOutputs1D_allnodes(xx[i,:],p).ravel()+vv[i,:]
    









t= list(range(Nsim+1))
# plt.figure(1)
# for i in range (Nz):
#     plt.subplot(5, 7, i+1)
#     plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b')
    

plt.figure(2)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    plt.plot(t,xx[:,i],'r')
plt.show()